﻿using System.IO;
using Microsoft.Xna.Framework;
using Microsoft.Xna.Framework.Graphics;
using System;
using System.Collections.Generic;
using System.Windows;
using System.Xml;

namespace Babylon.Toolbox
{
    public class Effect : IDisposable
    {
        VertexShader vertexShader;
        PixelShader pixelShader;
        GraphicsDevice device;
        readonly Dictionary<string, EffectParameter> parameters = new Dictionary<string, EffectParameter>();

        public GraphicsDevice Device
        {
            get { return device; }
            // added setter in order to be able to load the effect on the UI thread
            // and apply it on the Compositor thread
            set { device = value; }
        }

        public Effect(GraphicsDevice device, string assemblyName, string rootName)
            : this(device, assemblyName, rootName, rootName)
        {
        }

        public Effect(GraphicsDevice device, string assemblyName, string vertexRootName, string pixelRootName)
        {
            this.device = device;

            // Uri
            string vertexRootUri = string.Format(@"/{0};component/{1}.vs", assemblyName, vertexRootName);
            string pixelRootUri = string.Format(@"/{0};component/{1}.ps", assemblyName, pixelRootName);

            // Shaders
            Stream vertexShaderStream = Application.GetResourceStream(new Uri(vertexRootUri, UriKind.Relative)).Stream;
            Stream pixelShaderStream = Application.GetResourceStream(new Uri(pixelRootUri, UriKind.Relative)).Stream;

            vertexShader = EffectCacheSystem.CompileVertexShader(Device, new Uri(vertexRootUri, UriKind.Relative));
            pixelShader = EffectCacheSystem.CompilePixelShader(Device, new Uri(pixelRootUri, UriKind.Relative));

            vertexShaderStream.Dispose();
            pixelShaderStream.Dispose();

            // Assembly codes
            ExtractConstantsRegisters(vertexRootUri + ".constants", false);
            ExtractConstantsRegisters(pixelRootUri + ".constants", true);
        }

        void ExtractConstantsRegisters(string rootUri, bool isForPixelShader)
        {
            using (Stream stream = Application.GetResourceStream(new Uri(rootUri, UriKind.Relative)).Stream)
            {
                XmlReader xmlReader = XmlReader.Create(stream);

                while (xmlReader.ReadToFollowing("Constant"))
                {
                    xmlReader.ReadToDescendant("Name");
                    string name = xmlReader.ReadElementContentAsString();

                    EffectParameter effectParameter;
                    if (parameters.ContainsKey(name))
                    {
                        effectParameter = parameters[name];
                    }
                    else
                    {
                        effectParameter = new EffectParameter(this, name);
                        parameters.Add(name, effectParameter);
                    }

                    xmlReader.ReadToNextSibling("RegisterIndex");
                    int index = xmlReader.ReadElementContentAsInt();

                    if (isForPixelShader)
                        effectParameter.PixelShaderRegisterIndex = index;
                    else
                        effectParameter.VertexShaderRegisterIndex = index;

                    xmlReader.ReadToNextSibling("RegisterCount");
                    effectParameter.RegisterCount = Math.Max(effectParameter.RegisterCount, xmlReader.ReadElementContentAsInt());
                }

                xmlReader.Close();
            }
        }

        public EffectParameter GetParameter(string name)
        {
            EffectParameter result;

            return !parameters.TryGetValue(name, out result) ? null : result;
        }

        // Shader constants
        public void SetConstant(string propertyName, Matrix matrix)
        {
            EffectParameter effectParameter = GetParameter(propertyName);

            if (effectParameter != null)
                effectParameter.SetValue(matrix);
        }

        public void SetConstant(string propertyName, Vector4 vector4)
        {
            EffectParameter effectParameter = GetParameter(propertyName);

            if (effectParameter != null)
                effectParameter.SetValue(vector4);
        }

        public void SetConstant(string propertyName, float float01, float float02, float float03, float float04)
        {
            EffectParameter effectParameter = GetParameter(propertyName);

            if (effectParameter != null)
                effectParameter.SetValue(float01, float02, float03, float04);
        }

        public void SetConstant(string propertyName, bool bool01, bool bool02, bool bool03, bool bool04)
        {
            EffectParameter effectParameter = GetParameter(propertyName);

            if (effectParameter != null)
                effectParameter.SetValue(bool01, bool02, bool03, bool04);
        }

        public void SetConstant(string propertyName, Color color, float alpha)
        {
            EffectParameter effectParameter = GetParameter(propertyName);

            if (effectParameter != null)
                effectParameter.SetValue(color, alpha);
        }

        public void SetConstant(string propertyName, Color color)
        {
            EffectParameter effectParameter = GetParameter(propertyName);

            if (effectParameter != null)
                effectParameter.SetValue(color);
        }

        public void SetConstant(string propertyName, Vector3 vector3)
        {
            EffectParameter effectParameter = GetParameter(propertyName);

            if (effectParameter != null)
                effectParameter.SetValue(vector3);
        }

        // Apply
        public virtual void Apply()
        {
            // Shaders
            if (vertexShader != EffectCacheSystem.GetVertexShaderCache(Device))
            {
                Device.SetVertexShader(vertexShader);
                EffectCacheSystem.SetVertexShaderCache(Device, vertexShader);
            }

            if (pixelShader != EffectCacheSystem.GetPixelShaderCache(Device))
            {
                Device.SetPixelShader(pixelShader);
                EffectCacheSystem.SetPixelShaderCache(Device, pixelShader);
            }

            // Parameters
            foreach (EffectParameter effectParameter in parameters.Values)
            {
                effectParameter.Apply();
            }
        }

        public virtual void Dispose()
        {
            vertexShader = null;
            pixelShader = null;
        }
    }
}
